#!/usr/bin/python3

"""
MD文档合并工具
"""

import os
import re
import sys
import argparse
from pathlib import Path

release_version = "1.0"
DIR_PATTERN = re.compile(r'^(\d+(?:\.\d+)*)-(.+)$')
ID_PATTERN = re.compile(r'^(\d+(?:\.\d+)*)')

def parse_arguments():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description="""Markdown文档合并工具
示例用法:
  python merge-all-markdowns.py -i ./input -o ./output/merged.md -v 1.0""",
        epilog="默认输出路径会自动创建不存在的目录"
    )

    parser.add_argument('-i', '--input',
                        default='../benchmarks',
                        help='输入目录路径 (默认: %(default)s)')
    parser.add_argument('-o', '--output',
                        default='../release/OpenCloudOS安全配置基线.md',
                        help='输出文件路径 (默认: %(default)s)')
    parser.add_argument('-v', '--version',
                        default='1.0',
                        help='输入文档版本（默认: %(default)s)）')
    return parser.parse_args()

def generate_main_title(output_path):
    filename = os.path.splitext(os.path.basename(output_path))[0]
    return filename.replace('_', ' ').replace('-', ' ') + f"v{release_version}"

def validate_paths(args):
    input_dir = os.path.abspath(args.input)
    output_path = os.path.abspath(args.output)

    if not os.path.isdir(input_dir):
        sys.exit(f"错误：输入目录不存在 [{input_dir}]")

    output_dir = os.path.dirname(output_path)
    os.makedirs(output_dir, exist_ok=True)

    return input_dir, output_path

def natural_sort_key(file_path):
    filename = Path(file_path).name
    id_match = ID_PATTERN.match(filename)
    if id_match:
        return tuple(map(int, id_match.group(1).split('.')))
    return ()

def collect_md_files(root_dir):
    md_files = []
    for dirpath, _, filenames in os.walk(root_dir):
        for filename in filenames:
            if filename.endswith('.md'):
                full_path = os.path.join(dirpath, filename)
                md_files.append(full_path)
    md_files.sort(key=natural_sort_key)
    return md_files

def adjust_heading_level(content, level_increase):
    adjusted = []
    for line in content.split('\n'):
        if line.startswith('#'):
            match = re.match(r'^(#+)\s+(.*)', line)
            if match:
                hashes, title = match.groups()
                new_hashes = '#' * (len(hashes) + level_increase)
                line = f'{new_hashes} {title}'
        adjusted.append(line)
    return '\n'.join(adjusted)

def merge_md_files(root_dir, output_file, main_title):
    processed_dirs = set()

    md_files = collect_md_files(root_dir)

    with open(output_file, 'w', encoding='utf-8') as out_f:
        # 写入文档主标题
        out_f.write(f"# {main_title}\n\n")

        for md_file in md_files:
            # 获取相对路径并拆解路径组件
            rel_path = os.path.relpath(md_file, root_dir)
            path_components = Path(rel_path).parts[:-1]  # 排除文件名

            current_dir = root_dir  # 当前处理的目录路径
            dir_level = 0  # 目录层级计数器

            # 遍历路径组件生成章节标题
            for comp in path_components:
                current_dir = os.path.join(current_dir, comp)
                dir_match = DIR_PATTERN.match(comp)

                if dir_match and current_dir not in processed_dirs:
                    # 解析目录信息
                    chapter_num, raw_title = dir_match.groups()
                    chapter_title = raw_title.replace('-', ' ')

                    # 计算标题层级（根目录为1级，每级子目录+1）
                    rel_dir = os.path.relpath(current_dir, root_dir)
                    chapter_depth = sum(1 for c in Path(rel_dir).parts if DIR_PATTERN.match(c))

                    # 生成对应层级的Markdown标题
                    heading = '#' * (chapter_depth + 1)  # +1因为主标题已占用第一级
                    heading_line = f"{heading} {chapter_num} {chapter_title}\n"

                    # 写入标题并记录已处理目录
                    out_f.write(heading_line + '\n')
                    processed_dirs.add(current_dir)

                dir_level += 1

            with open(md_file, 'r', encoding='utf-8') as in_f:
                content = in_f.read()

            rel_dir = os.path.dirname(rel_path)
            level_increase = sum(1 for c in Path(rel_dir).parts if DIR_PATTERN.match(c))

            adjusted_content = adjust_heading_level(content, level_increase)
            out_f.write(adjusted_content + '\n\n')

if __name__ == '__main__':
    args = parse_arguments()
    input_dir, output_path = validate_paths(args)
    release_version = args.version
    document_title = generate_main_title(output_path)

    try:
        merge_md_files(
            root_dir=input_dir,
            output_file=output_path,
            main_title=document_title
        )
        print(f"合并成功！生成文件：{output_path}")
    except Exception as e:
        sys.exit(f"处理过程中发生错误：{str(e)}")
